#include "libcomm.h"

Process::Process()
{
#ifdef COMM_MSVC
	ZeroMemory( &m_pi, sizeof(m_pi) );
#else
	m_child = -1;
#endif
}

Process::~Process()
{
#ifdef COMM_MSVC
	if(NULL != m_pi.hProcess)
	{
		CloseHandle( m_pi.hProcess );
		m_pi.hProcess = NULL;
	}
	if(NULL != m_pi.hThread)
	{
		CloseHandle( m_pi.hThread );
		m_pi.hThread = NULL;
	}
#endif
}

void Process::Terminate()
{
#ifdef COMM_MSVC
	if(NULL != m_pi.hProcess)
	{
		if(!::TerminateProcess(m_pi.hProcess,-1))
			LOGFUNCERROR(TerminateProcess);
	}
	
#else
	if(-1 != m_child)
	{
		if(-1 == ::kill(m_child,SIGKILL))
			LOGFUNCERROR(kill);
	}
	
#endif
}

const char* Process::Name() const
{
	return m_filename;
}

bool Process::Join(int* exitcode)
{
#ifdef COMM_MSVC

	DWORD nExitCode;
	DWORD dwEvent;
	MYLIB_ASSERT_NOTNULL(m_pi.hProcess);

	LOGNOTICE("waiting for process \"%s\" exit...", (char*)m_filename);
	dwEvent = WaitForSingleObject( m_pi.hProcess, INFINITE );
	if(WAIT_OBJECT_0 != dwEvent)
	{
		LOGFUNCERROR(WaitForSingleObject);
		return false;
	}

	if(!GetExitCodeProcess(m_pi.hProcess,&nExitCode))
	{
		LOGFUNCERROR(GetExitCodeProcess);
		return false;
	}

	LOGNOTICE("process \"%s\" exited: %d.", (char*)m_filename, nExitCode);
	if(NULL != exitcode)
		*exitcode = nExitCode;
	return true;

#else

	pid_t pw;
	int status;
	MYLIB_ASSERT(-1 != m_child);
	do 
	{
		LOGNOTICE("waiting for process \"%s\" exit...", (char*)m_filename);
		pw = waitpid(m_child,&status,WUNTRACED|WCONTINUED);
		if (-1 == pw) 
		{
			LOGFUNCERROR(waitpid);
			break;
		}

		if (WIFEXITED(status))
		{
			LOGNOTICE("process \"%s\" exited: %d.", (char*)m_filename, WEXITSTATUS(status));
			if(NULL != exitcode)
				*exitcode = WEXITSTATUS(status);
			return true;
		} 

		if (WIFSIGNALED(status)) {
			LOGWARN("process \"%s\" killed by signal %d.", (char*)m_filename, WTERMSIG(status));
		} else if (WIFSTOPPED(status)) {
			LOGWARN("process \"%s\" stopped by signal %d.", (char*)m_filename, WSTOPSIG(status));
		} else if (WIFCONTINUED(status)) {
			LOGWARN("process \"%s\" continued.", (char*)m_filename);
		}

	} while (!WIFEXITED(status) && !WIFSIGNALED(status));
	return false;

#endif
}

Process* Process::Execute(const char* filename, const Arguments& argv, const Arguments& envp, bool hide)
{
	Process* process = NULL;
	
#ifdef COMM_MSVC

	String cmd = String::format("\"%s\"",filename);
	int argc = MYLIB_MIN(PROCESS_MAX_ARGS,argv.count());
	for (int i=0;i<argc;i++)
		cmd << " " << argv[i].ToString();

	do 
	{
		LOGNOTICE("executing process \"%s\"...",filename);

		STARTUPINFOA si;
		PROCESS_INFORMATION pi;
		ZeroMemory( &si, sizeof(si) );
		si.cb = sizeof(si);
		si.dwFlags = STARTF_USESHOWWINDOW;
		si.wShowWindow = hide?SW_HIDE:SW_SHOW;
		ZeroMemory( &pi, sizeof(pi) );

		// Start the child process. 
		if( !CreateProcessA( NULL, // No module name (use command line). 
			cmd,			  // Command line. 
			NULL,             // Process handle not inheritable. 
			NULL,             // Thread handle not inheritable. 
			FALSE,            // Set handle inheritance to FALSE. 
			0,                // No creation flags. 
			NULL,             // Use parent's environment block. 
			NULL,             // Use parent's starting directory. 
			&si,              // Pointer to STARTUPINFO structure.
			&pi )             // Pointer to PROCESS_INFORMATION structure.
			) 
		{
			LOGFUNCERROR(CreateProcessA);
			break;
		}

		process = new Process();
		MYLIB_ASSERT_NOTNULL(process);
		process->m_pi = pi;
		process->m_filename = filename;

	} while (0);


#else

	int argc = 0;
	int i = 0;
	StringList argvList = argv.ToStringList();
	StringList envpList = envp.ToStringList();
	char* argv_all[PROCESS_MAX_ARGS+2];
	char* envp_all[PROCESS_MAX_ARGS+1];
	argvList.insert(0,filename);

	argc = MYLIB_MIN(PROCESS_MAX_ARGS+1,argvList.count());
	for (i=0;i<argc;i++)
		argv_all[i] = argvList[i].get();
	argv_all[i] = NULL;

	argc = MYLIB_MIN(PROCESS_MAX_ARGS,envpList.count());
	for (i=0;i<argc;i++)
		envp_all[i] = envpList[i].get();
	envp_all[i] = NULL;

	do 
	{
		LOGNOTICE("executing process \"%s\"...",filename);
		if(NULL != strrchr(filename,'/') && -1 == access(filename,X_OK))
		{
			LOGWARN("executing process \"%s\" failure: %s.",filename,strerror(errno));
			break;
		}

		pid_t child = fork();
		switch (child) 
		{
		case -1:  
			{
				LOGFUNCERROR(fork);
				break;
			}
		case 0:
			{
			if (NULL == strrchr(filename, '/'))
#ifdef __CYGWIN__
				#warning "!!!envionment variables was not supported in cygwin version!!!"
				execvp(filename, argv_all);
#else
				execvpe(filename, argv_all, envp_all);
#endif
				else
#ifdef __CYGWIN__
					#warning "!!!envionment variables was not supported in cygwin version!!!"
					execv(filename,argv_all);
#else
					execve(filename, argv_all, envp_all);
#endif
				LOGFUNCERROR(execve);
				exit(-1);
			}
		default:            
			{
				process = new Process();
				MYLIB_ASSERT_NOTNULL(process);
				process->m_child = child;
				process->m_filename = filename;
				break;
			}
		}//switch
	} while (0);

#endif

	return process;
}

bool Process::ExecuteSync(const char* filename, int* exitcode, const Arguments& argv, const Arguments& envp, bool hide)
{
	bool ret = false;
	Process* process = Execute(filename,argv,envp,hide);
	if(process != NULL)
	{
		ret = process->Join(exitcode);
		delete process;
	}
	return ret;
}


bool Process::ExecuteSystemCommand(const char* command, int* exitcode)
{
	MYLIB_ASSERT_NOTNULL(command);
	LOGNOTICE("executing system command \"%s\"...",command);
	int status = system(command);
#ifdef COMM_MSVC

	if(-1 == status)
	{
		LOGFUNCERROR(system);
		return false;
	}

	if(NULL != exitcode)
		*exitcode = status;
	return true;

#else

	if(-1 == status)
	{
		LOGERR("Function \"system\" error: child process could not be created, or its status could not be retrieved.");
		return false;
	}
	if(127 == status)
	{
		LOGERR("Function \"system\" error: shell could not be executed in the child process.");
		return false;
	}

	if (WIFEXITED(status))
	{
		LOGNOTICE("system command exited: %d.", WEXITSTATUS(status));
		if(NULL != exitcode)
			*exitcode = WEXITSTATUS(status);
		return true;
	} 

	if (WIFSIGNALED(status)) {
		LOGWARN("system command killed by signal %d.", WTERMSIG(status));
	} else if (WIFSTOPPED(status)) {
		LOGWARN("system command stopped by signal %d.", WSTOPSIG(status));
	} else if (WIFCONTINUED(status)) {
		LOGWARN("system command continued.");
	}

	return false;

#endif
}

#ifndef COMM_MSVC
StringList Process::EchoSystemCommand(const char* command,int maxLineSize)
{
	StringList lines;
	FILE *pp = popen(command, "r");
	if(NULL != pp)
	{
		Buffer tmp(maxLineSize);
		while (fgets(tmp, tmp.size()-1, pp) != NULL) 
		{
			lines.append(tmp.get());
		}
		pclose(pp); 
	}
	else
	{
		LOGFUNCERROR(popen);
	}
	return lines;
}
#endif


#ifdef COMM_MSVC
#include <UserEnv.h>
#include <WtsApi32.h>
#include <ShellAPI.h>
#pragma comment(lib, "Wtsapi32.lib")
#pragma comment(lib, "Userenv.lib")
Process* Process::ExecuteAsUser(const char* filename, const Arguments& argv, int timeoutMsForUserLogin, bool hide)
{
	Process* process = NULL;
	String cmd = String::format("\"%s\"",filename);
	int argc = min(PROCESS_MAX_ARGS,argv.count());
	for (int i=0;i<argc;i++)
		cmd << " " << argv[i].ToString();

	HANDLE hToken = NULL;
	HANDLE hTokenDup = NULL;
	LPVOID pEnv = NULL;
	do 
	{
		LOGNOTICE("executing user process \"%s\"...",filename);

		STARTUPINFOA si;
		PROCESS_INFORMATION pi;
		do 
		{
			DWORD dwSessionID = WTSGetActiveConsoleSessionId();
			if (!WTSQueryUserToken(dwSessionID, &hToken))
			{
				LOGFUNCERROR(WTSQueryUserToken);
				if(timeoutMsForUserLogin)
				{
					LOGWARN("waiting(%dms) for user login...",timeoutMsForUserLogin);
					Thread::Sleep(timeoutMsForUserLogin);
					continue;
				}
			}
			break;
		} while (1);
		if(NULL == hToken)
			break;

		if (!DuplicateTokenEx(hToken, MAXIMUM_ALLOWED, NULL, SecurityIdentification, TokenPrimary, &hTokenDup))
		{
			LOGFUNCERROR(DuplicateTokenEx);
			break;
		}

		if (!CreateEnvironmentBlock(&pEnv, hTokenDup, FALSE))
		{
			LOGFUNCERROR(CreateEnvironmentBlock);
			break;
		}

		ZeroMemory( &si, sizeof(si) );
		si.cb = sizeof(si);
		si.dwFlags = STARTF_USESHOWWINDOW;
		si.wShowWindow = hide?SW_HIDE:SW_SHOW;
		ZeroMemory( &pi, sizeof(pi) );

		DWORD dwCreateFlag = NORMAL_PRIORITY_CLASS | CREATE_NEW_CONSOLE | CREATE_UNICODE_ENVIRONMENT;
		if (!CreateProcessAsUserA(hTokenDup, NULL, cmd, NULL, NULL, FALSE, dwCreateFlag, pEnv, NULL, &si, &pi))
		{
			LOGFUNCERROR(CreateProcessAsUserA);
			break;
		}

		process = new Process();
		MYLIB_ASSERT_NOTNULL(process);
		process->m_pi = pi;
		process->m_filename = filename;

	} while(0);

	if(NULL != pEnv)
		DestroyEnvironmentBlock(pEnv);
	if(NULL != hTokenDup)
		CloseHandle(hTokenDup);
	if(NULL != hToken)
		CloseHandle(hToken);

	return process;
}

Process* Process::ExecuteAsAdmin(const char* filename, const Arguments& argv, bool hide)
{
	Process* process = NULL;
	String cmd = String::format("\"%s\"",filename);
	int argc = min(PROCESS_MAX_ARGS,argv.count());
	for (int i=0;i<argc;i++)
		cmd << " " << argv[i].ToString();

	do 
	{
		LOGNOTICE("executing admin process \"%s\"...",filename);

		SHELLEXECUTEINFOA ShExecInfo = {0};  
		ShExecInfo.cbSize = sizeof(SHELLEXECUTEINFO);  
		ShExecInfo.fMask = SEE_MASK_NOCLOSEPROCESS;  
		ShExecInfo.hwnd = NULL;  
		ShExecInfo.lpVerb = "runas";  
		ShExecInfo.lpFile = filename;  
		ShExecInfo.lpParameters = (const char*)cmd;
		ShExecInfo.lpDirectory = NULL;  
		ShExecInfo.nShow = hide?SW_HIDE:SW_SHOW;  
		ShExecInfo.hInstApp = NULL;  
		if(!ShellExecuteExA(&ShExecInfo))
		{
			LOGFUNCERROR(ShellExecuteExA);
			break;
		}

		process = new Process();
		MYLIB_ASSERT_NOTNULL(process);
		process->m_pi.hProcess = ShExecInfo.hProcess;
		process->m_pi.hThread = NULL;
		process->m_filename = filename;

	} while (0);

	return process;
}

bool Process::IsAdminPrivilege()
{
	UINT16 uWinVer = LOWORD(GetVersion());
	uWinVer = MAKEWORD(HIBYTE(uWinVer),LOBYTE(uWinVer));
	if (uWinVer < 0x0600)
		return true;

	SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
	PSID AdministratorsGroup;

	if(!AllocateAndInitializeSid(
		&NtAuthority,
		2,
		SECURITY_BUILTIN_DOMAIN_RID,
		DOMAIN_ALIAS_RID_ADMINS,
		0, 0, 0, 0, 0, 0,
		&AdministratorsGroup))
	{
		LOGFUNCERROR(AllocateAndInitializeSid);
		return false;
	}

	BOOL isAdm = FALSE;
	if(!CheckTokenMembership(NULL, AdministratorsGroup, &isAdm))
		LOGFUNCERROR(CheckTokenMembership);
	FreeSid(AdministratorsGroup);
	return (isAdm?true:false);
}
#endif



